import copy
import logging
import os
import re
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path

from tqdm import tqdm

from gitmanager.projectstate import GitProject
from llm.client import DeepSeekClient, APIMode
from parser import ProjectAnalyzer, CProjectLoader
from parser.graphfunc import get_adjacency
from parser.rustfunc import remove_rust_comments
from prompts import TranslatePrompt, FixPrompt
from .fixer import RustFixer
from utils import system_call, write_file, find_most_similar, find_most_similar_uncased, split_by_non_var_chars
from utils.progress import ProgressServer

TRANSLATE_PROMPT_VERSION = 'version1'
FIX_PROMPT_VERSION = 'version1'
DEBUG = False
TRANSLATE_DISABLED = False
LOG_PATH = Path(os.environ['LOG_PATH'])

def _match_name(name, ds):
    found_decl = None
    for d in ds:
        if d.name == name:
            found_decl = d
    return found_decl

def match_decl_name(name, all_decls):
    all_decls = list(all_decls)
    found_decl = _match_name(name, all_decls)
    if found_decl is None:
        similar_name = find_most_similar(name, map(lambda x: x.name, all_decls))
        found_decl = _match_name(similar_name, all_decls)
    return found_decl


class Translator:
    _progress_bar = None

    def __init__(self,
                 c_project_path: Path | str,
                 rust_project_path: Path | str,
                 project_type='lib',
                 model='',
                 fix_rounds=5,
                 parallel_file=2,
                 parallel_unit=5,
                 temp_path='./temp'):
        self.temp_path = Path(temp_path).resolve()

        self.c_project_path = Path(c_project_path).resolve()
        self.rust_project_path = Path(rust_project_path).resolve()
        self.project_type = project_type
        self.model = model
        self.fix_rounds = fix_rounds
        self.parallel_file = parallel_file
        self.parallel_unit = parallel_unit

        self.file_unit_info_map = None

    @staticmethod
    def to_rust_name(name):
        """
        将名称转换为符合 Rust 项目命名规则的 snake_case 格式。

        :param name: str, 原始名称
        :return: str, 转换后的名称
        """
        # 去除多余的空格
        name = name.strip()

        # 将 PascalCase 或 camelCase 转为 snake_case
        name = re.sub(r'(?<!^)(?=[A-Z])', '_', name).lower()

        # 替换其他非法字符为下划线（仅保留字母、数字和下划线）
        name = re.sub(r'[^a-z0-9_]', '_', name)

        # 移除多余的连续下划线
        name = re.sub(r'_+', '_', name)

        # 如果项目名以数字开头，添加前缀
        m = re.match(r'^(\d*)', name)
        if m:
            name = name.replace(m.group(1), '')

        # 确保不以下划线开头或结尾
        name = name.strip('_')

        return name

    def _prepare(self):
        self.c_project = ProjectAnalyzer(self.c_project_path)
        self._init_rust()
        self._init_rust_code_file()
        self._init_rust_code_import()
        self.git_rust_project = GitProject.git_init(self.rust_project_path, temp_path=self.temp_path)

    def _init_rust_code_import(self):
        file_graph = self.c_project.file_dependency.origin_graph
        get_rust_code_path = lambda m: self.module_map[m].relative_to(self.rust_project_path)
        for module, data in file_graph.nodes(data=True):
            import_decls = []
            rust_code_path = get_rust_code_path(module)
            for depend in get_adjacency(file_graph, module):
                depend_rust_code_path = get_rust_code_path(depend)
                if depend_rust_code_path.is_relative_to('tests') and rust_code_path.is_relative_to('src'):
                    import_decl = (f"#[cfg(test)]\n"
                                   f"#[path = \"../{depend_rust_code_path}\"]\n"
                                   f"mod {depend_rust_code_path.stem};\n\n"
                                   f"#[cfg(test)]\n"
                                   f"use {depend_rust_code_path.stem}::*;")
                    import_decls.append(import_decl)
                elif depend_rust_code_path.is_relative_to('tests') and rust_code_path.is_relative_to('tests'):
                    import_decl = (f"mod {depend_rust_code_path.stem};\n\n"
                                   f"use crate::{depend_rust_code_path.stem}::*;")
                    import_decls.append(import_decl)
                else:
                    suffix = depend_rust_code_path.suffix
                    crate = "::".join(map(self.to_rust_name,
                                          Path(str(depend_rust_code_path).removesuffix(suffix)).parts[1:]))
                    if crate:
                        if rust_code_path.is_relative_to('tests') or rust_code_path == Path('src/main.rs'):
                            import_decl = f"use {self.project_name}::{crate}::*;"
                        else:
                            import_decl = f"use crate::{crate}::*;"
                        import_decls.append(import_decl)
            decls_content = '\n'.join(import_decls)
            if decls_content:
                decls_content += '\n'
            write_file(self.module_map[module], decls_content)

    def _init_rust_code_file(self):
        file_levels, _ = self.c_project.get_parallel_levels()
        nodes = []
        for level in file_levels:
            for node in level:
                if isinstance(node, list):
                    for n in node:
                        module, info = n
                        nodes.append((module, info))
                else:
                    module, info = node
                    nodes.append((module, info))
        self.module_map = {}
        for module, info in nodes:
            file = info['c_file']
            if file is None:
                file = info['h_file']
            file = Path(file)
            rust_code_file = self._add_rust_code_file(file)
            self.module_map[module] = self.rust_project_path / rust_code_file

    def _c_file_trans_rust_file(self, c_file):
        c_file = Path(c_file)
        c_file = c_file.relative_to(self.c_project_path)
        parts = list(Path(c_file).parts)

        name = parts[-1]
        name = Path(name).stem
        name = self.to_rust_name(name)

        if parts[0] == 'test':
            parts[0] = 'tests'

        rust_code_path = Path(*parts[:-1], f'{name}.rs')
        return rust_code_path

    def _add_rust_code_file(self, c_file):
        rust_code_path = self._c_file_trans_rust_file(c_file)
        self._add_path_mod_declaration(rust_code_path)
        write_file(self.rust_project_path / rust_code_path, '')
        return rust_code_path

    def _add_path_mod_declaration(self, rust_code_path: Path):
        mods = list(rust_code_path.parts)
        current_path = self.rust_project_path / mods[0]
        for i in range(1, len(mods) - 1):
            mod = mods[i]
            temp_dir = current_path / mod
            temp_dir.mkdir(parents=True, exist_ok=True)
            declaration_file = self.decision_declaration_file(mods[0], i)
            self.add_mod_declaration(current_path, declaration_file, mod)
            current_path = current_path / mod
        mod = mods[-1]
        mod = Path(mod).stem
        declaration_file = self.decision_declaration_file(mods[0], len(mods) - 1)
        self.add_mod_declaration(current_path, declaration_file, mod)

    @staticmethod
    def decision_declaration_file(main_dirname, dir_level):
        if dir_level > 1:
            return 'mod.rs'
        elif dir_level == 1:
            if main_dirname == 'src':
                return 'lib.rs'
            elif main_dirname == 'tests':
                return 'tests.rs'
        raise ValueError(f'不符合的输入: (main_dirname, dir_level) -> {main_dirname}, {dir_level}')

    @staticmethod
    def add_mod_declaration(path, declaration_file, declaration_mod):
        if declaration_file == 'lib.rs' and declaration_mod == 'main':
            return
        with open(path / declaration_file, 'a', encoding='utf-8') as fp:
            fp.write(f'pub mod {declaration_mod};\n')

    def _init_rust(self):
        # 初始化rust项目
        self.project_name = self.c_project_path.name
        self.project_name = self.to_rust_name(self.project_name)

        self.rust_project_path.mkdir(parents=True, exist_ok=True)

        current_dir = os.getcwd()
        os.chdir(self.rust_project_path)
        if self.project_type == 'lib':
            system_call(f'cargo init --name {self.project_name} --lib')
            with open(os.path.join('src', 'lib.rs'), 'w', encoding='utf-8'):
                pass
        else:
            system_call(f'cargo init --name {self.project_name} --bin')
            with open(os.path.join('src', 'main.rs'), 'w', encoding='utf-8'):
                pass
        os.chdir(current_dir)
        self.rust_project_test = self.rust_project_path / 'tests'
        self.rust_project_test.mkdir(parents=True, exist_ok=True)

    def _create_module_task(self, module, module_levels, module_map):
        parallel_level = module_levels[module]
        rust_file = self.module_map[module].relative_to(self.rust_project_path)

        unit_graph = self.c_project.function_dependencies[module].origin_graph.copy()

        task = (self.git_rust_project, self.c_project.project,
                module, rust_file,
                parallel_level,
                unit_graph, module_map,
                self.parallel_unit)
        return task

    @staticmethod
    def _calculate_module_workload(module, module_levels, counter):
        parallel_level = module_levels[module]
        for level in parallel_level:
            for unit in level:
                if isinstance(unit, list):
                    for _ in unit:
                        counter[0] += 1
                else:
                    counter[0] += 1

    def _calculate_total_workload(self, file_levels, module_levels):
        counter = [0]
        for level in file_levels:
            for node in level:
                if isinstance(node, list):
                    for n in node:
                        module, info = n
                        self._calculate_module_workload(module, module_levels, counter)
                else:
                    module, info = node
                    self._calculate_module_workload(module, module_levels, counter)
        return counter[0]

    def translate(self):
        self._prepare()
        file_levels, module_levels = self.c_project.get_parallel_levels()
        file_unit_info_map = self.c_project.get_file_unit_info_map()
        self.file_unit_info_map = file_unit_info_map

        total_work = self._calculate_total_workload(file_levels, module_levels)

        progress_server = None
        if DEBUG:
            Translator._progress_bar = tqdm(total=total_work, desc='Debugging')
        else:
            progress_server = ProgressServer(total_work, 'Translating')

        if not DEBUG:
            progress_server.run_server()

        for level in file_levels:
            # 传git，C项目代码，目标文件名，解算图，目标文件函数依赖图，全部的函数依赖图(参考用, 会保存翻译后的rust声明)
            tasks = []
            for node in level:
                if isinstance(node, list):
                    group_task = []
                    for n in node:
                        module, info = n
                        task = self._create_module_task(module, module_levels, file_unit_info_map)
                        group_task.append(task)
                    tasks.append(group_task)
                else:
                    module, info = node
                    task = self._create_module_task(module, module_levels, file_unit_info_map)
                    tasks.append(task)
            if DEBUG:
                result = [self._translate_module_task(task) for task in tasks]
            else:
                with ProcessPoolExecutor(max_workers=self.parallel_file) as executor:
                    # 返回(git, {module: unit_map})
                    result = list(executor.map(self._translate_module_task, tasks))
            for status, module_map in result:
                self.git_rust_project.merge(status)
                for module, unit_map in module_map.items():
                    file_unit_info_map[module] = unit_map
                status.remove()

        if not DEBUG:
            progress_server.stop_server()
        else:
            Translator._progress_bar.close()

        # self.git_rust_project.visualize(self.temp_path / 'git_graph')
        # self.git_rust_project.repo.close()

    def fix(self):
        logger = logging.getLogger()

        fix_prompt = FixPrompt(FIX_PROMPT_VERSION).fix
        client = DeepSeekClient(os.environ['API_KEY'])

        fixer = RustFixer(self.git_rust_project, client, fix_prompt, LOG_PATH, self.fix_rounds, progress=True)

        if not fixer.fix_rust_code():
            logger.error(f"修复后仍然失败")

    def run(self):
        self.translate()
        # self.fix()

        self.git_rust_project.visualize(self.temp_path / 'git_graph')
        self.git_rust_project.repo.close()

    @staticmethod
    def _translate_module_task(task):
        if isinstance(task, list):
            status, module_map = (
                Translator._translate_module_group(task)
            )
        else:
            status, module_map = (
                Translator._translate_module(*task)
            )
        return status, module_map

    @staticmethod
    def _translate_module_group(group_task):
        git_project = group_task[0][0]
        modules = map(lambda x: x[2], group_task)
        branch_name = ','.join(modules)
        temp_project = git_project.create_branch(branch_name)
        return_module_map = {}

        for task in group_task:
            status, module_map = Translator._translate_module(temp_project, *task[1:])
            temp_project.merge(status)
            return_module_map.update(module_map)
            status.remove()
        return temp_project, return_module_map

    @staticmethod
    def _create_unit_task(git_project, c_project_code,
                          unit, unit_info,
                          module, rust_file,
                          unit_graph, module_map):
        task = (git_project, c_project_code,
                unit, unit_info,
                module, rust_file,
                unit_graph, module_map)
        return task

    @staticmethod
    def _translate_module(git_project, c_project_code,
                          module, rust_file,
                          parallel_level,
                          unit_graph, module_map,
                          parallel_unit):
        temp_project = git_project.create_branch(module)
        current_module_unit_map = copy.deepcopy(module_map[module])

        for level in parallel_level:
            # 传 git，C项目代码，unit, unit_info, 目标文件名，目标文件函数依赖图(要更新)，目标文件的依赖文件的函数依赖图(参考用)
            tasks = []
            for unit in level:
                if isinstance(unit, list):
                    group_task = []
                    for u, info in unit:
                        task = Translator._create_unit_task(temp_project, c_project_code,
                                                            u, info,
                                                            module, rust_file,
                                                            unit_graph, module_map)
                        group_task.append(task)
                    tasks.append(group_task)
                else:
                    u, info = unit
                    task = Translator._create_unit_task(temp_project, c_project_code,
                                                        u, info,
                                                        module, rust_file,
                                                        unit_graph, module_map)
                    tasks.append(task)
            if DEBUG:
                result = [Translator._translate_unit_task(task) for task in tasks]
            else:
                with ProcessPoolExecutor(max_workers=parallel_unit) as executor:
                    # 返回 (git, 声明: [unit, rust_unit_decl])
                    result = list(executor.map(Translator._translate_unit_task, tasks))
            for status, decls in result:
                if status:
                    temp_project.merge(status)
                    for decl in decls:
                        unit, rust_decl = decl
                        current_module_unit_map[unit]['rust_decl'] = rust_decl
                    status.remove()
        return temp_project, {module: current_module_unit_map}

    @staticmethod
    def _translate_unit_task(task):
        if isinstance(task, list):
            status, decls = (
                Translator._translate_unit_group(task)
            )
        else:
            status, decls = (
                Translator._translate_unit(*task)
            )
        return status, decls

    @staticmethod
    def _translate_unit_group(group_task):
        if all(map(lambda x: x[3]['extern'], group_task)):
            if DEBUG:
                Translator._progress_bar.update(len(group_task))
            else:
                ProgressServer.update_progress(len(group_task))
            return None, []

        git_project = group_task[0][0]
        units = []
        for unit in map(lambda x: x[2], group_task):
            unit = unit.split(' ')[-1]
            units.append(unit)
        units_name = ','.join(units)
        module = group_task[0][4]

        temp_project = git_project.create_branch(f'{module}.{units_name}')
        return_decls = []

        for task in group_task:
            status, decls = Translator._translate_unit(temp_project, *task[1:])
            if status:
                temp_project.merge(status)
                return_decls.extend(decls)
                status.remove()
        return temp_project, return_decls

    @staticmethod
    def _translate_unit(git_project: GitProject, c_project_code: CProjectLoader,
                        unit, unit_info,
                        module, rust_file: Path,
                        unit_graph, module_map):
        unit_extern = unit_info['extern']
        if unit_extern:
            if DEBUG:
                Translator._progress_bar.update()
            else:
                ProgressServer.update_progress(1)
            return None, []

        logger = logging.getLogger()

        unit_name = unit.split(' ')[-1]
        temp_project = git_project.create_branch(f'{module}.{unit_name}')

        client = DeepSeekClient(os.environ['API_KEY'])

        rust_file = temp_project.source_path / rust_file
        # struct function global
        unit_type = unit_info['type']

        prompts = TranslatePrompt(TRANSLATE_PROMPT_VERSION)
        prompt = prompts[unit_type]

        def get_code(location: str):
            _path, start, end = location.split(':')
            start, end = int(start), int(end)
            return ''.join(c_project_code.get_file_lines(_path, start, end))

        if TRANSLATE_DISABLED:
            with open(rust_file, 'a', encoding='utf-8') as file:
                file.write(get_code(unit_info['location']))

            temp_project.commit([rust_file],
                                f'{unit}: {unit_info['location']}\n'
                                f'{unit}: {rust_file}')

            if DEBUG:
                Translator._progress_bar.update()
            else:
                ProgressServer.update_progress(1)

            return temp_project, [(unit, '')]

        unit_dependencies = get_adjacency(unit_graph, unit)
        unit_map = module_map[module]

        extern_units = {}
        static_units = {}
        for u, u_info in unit_dependencies.items():
            code = get_code(u_info['location'])
            if u_info['extern']:
                path = Path(u_info['location'].split(':')[0])
                extern_module = path.stem
                unit_map_info = copy.deepcopy(module_map[extern_module][u])
                unit_map_info['extern_module'] = extern_module
                unit_map_info['code'] = code
                extern_units[u] = unit_map_info
            else:
                unit_map_info = copy.deepcopy(unit_map[u])
                unit_map_info['code'] = code
                static_units[u] = unit_map_info

        c_dependency_code = "\n".join([
            # 内部依赖（当前unit的unit_graph的邻接结点）
            *[u_info['code']
              for u, u_info in static_units.items()],

            # 外部依赖（当前unit的unit_graph的邻接结点的extern部分）
            *[u_info['code']
              for u, u_info in extern_units.items()]
        ])
        # module map 才会更新 rust_decl，graph里面的结点更新没写
        translated_rust_code = "\n".join([
            # 内部依赖的翻译结果（当前rust文件）
            *[u_info['rust_decl']
              for u, u_info in static_units.items()
              if 'rust_decl' in u_info],

            # 外部依赖的翻译结果（其他模块）
            *[u_info['rust_decl']
              for u, u_info in extern_units.items()
              if 'rust_decl' in u_info]
        ])

        # 当前单元原始C代码
        c_code_to_translate = get_code(unit_info['location'])

        actual_prompt = prompt(c_code_to_translate=c_code_to_translate,
                               c_dependency_code=c_dependency_code,
                               translated_rust_code=translated_rust_code,
                               unit=unit)

        response = client.generate(actual_prompt, mode=APIMode.COMPLETION)

        # response = "```rust```"

        m = re.search(r"```rust\n(.*?)```", response, re.DOTALL)
        if m:
            rust_code = m.group(1)
            with open(rust_file, 'a', encoding='utf-8') as file:
                file.write(rust_code)
        else:
            logger.error(f"翻译失败, 代码: {module}/{unit}")
            client.save_histories(LOG_PATH / module / f'{unit}.jsonl')
            return temp_project, [(unit, '')]

        # 搜索到rust_decl只用来写commit_message
        clean_rust_code = remove_rust_comments(rust_code)
        clean_rust_code = clean_rust_code.splitlines()
        clean_rust_code = filter(lambda x: len(x) > 0, clean_rust_code)
        clean_rust_code = list(clean_rust_code)
        decl_words = split_by_non_var_chars(clean_rust_code[0])
        match_words = list(filter(lambda x: x == unit, decl_words))
        if match_words:
            rust_decl = match_words[0]
        else:
            rust_decl = find_most_similar_uncased(unit, decl_words)[0]

        # 找到翻译后的声明
        # rust_decl = None
        # # 这是以前使用解析器去找rust_code
        # parser = RustParser(rust_code)
        # declarations = parser.parse_all()
        #
        # # struct function global
        # if unit_type == 'function':
        #     decls = filter(lambda d: d.kind == 'function', declarations)
        #     decls = list(decls)
        #     if decls:
        #         rust_decl = match_decl_name(unit_name, decls)
        # elif unit_type == 'struct':
        #     decls = filter(lambda d: d.kind in {'struct', 'union', 'enum'}, declarations)
        #     decls = list(decls)
        #     if decls:
        #         rust_decl = match_decl_name(unit_name, decls)
        # elif unit_type == 'global':
        #     decls = filter(lambda d: d.kind == 'static', declarations)
        #     decls = list(decls)
        #     if decls:
        #         rust_decl = match_decl_name(unit_name, decls)
        # if rust_decl is None:
        #     logger.warning(f"没有找到对应的翻译声明, {module}/{unit}")

        fix_prompt = FixPrompt(FIX_PROMPT_VERSION).fix

        fixer = RustFixer(temp_project, client, fix_prompt, LOG_PATH)
        if not fixer.fix_rust_code(rust_file):
            logger.error(f"修复后仍然失败, 代码: {temp_project.source_path}: {module}/{unit}")

        client.save_histories(LOG_PATH / module / f'{unit}.jsonl')

        temp_project.commit([rust_file],
                            f'{unit}: {unit_info['location']}\n'
                            f'{rust_decl}: {rust_file}')

        # 根据commit找到代码
        rust_code = temp_project.get_current_commit_add_code(rust_file)

        if DEBUG:
            Translator._progress_bar.update()
        else:
            ProgressServer.update_progress(1)

        return temp_project, [(unit, rust_code)]
